import json

import urllib3

API_KEY = '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'
MAX_TOKENS = 4000
POOL = urllib3.PoolManager(5)
API_URL = 'https://api.openai.com/v1/chat/completions'#'https://api.openai.com/v1/completions'
API_HEADER = {
    'Authorization': f'Bearer {API_KEY}',
    'Content-Type': 'application/json;charset=UTF-8'
}
MODEL = 'gpt-3.5-turbo'#'text-davinci-003'

class Api:
    """
    获取答案api接口实现
    """

    @staticmethod
    def get_answer(question):
        """
        获取答案
        :param question: 问题
        :return: 答案
        """
        data = Api.get_answer_of_question_raw(question)
        data = Api.format_answer(data)
        # data = data['choices'][0]['text'].strip()
        return data

    @staticmethod
    def format_answer(answer):
        """
        格式化答案获取目标内容
        :param answer: 答案json
        :return: 答案字符串
        """
        data = json.loads(answer)
        if 'error' in data:
            raise Exception(data['error']['message'])
        data = data['choices'][0]['text'].strip()
        return data

    @staticmethod
    def get_answer_of_question_raw(question):
        """
        :param 问题，从列表中读取到
        :return: 该问题请求返回的json
        """
        data = {
            'prompt': question,
            'max_tokens': MAX_TOKENS,
            'temperature': 1,
            'model': MODEL,
        }
        request = POOL.urlopen(
            'POST',
            API_URL,
            body=json.dumps(data).encode('utf-8'),
            headers=API_HEADER
        )
        return request.data
